from rich import print

from collections import deque
from typing import Generator
from copy import deepcopy
from dataclasses import dataclass

from solver.SBoard import SBoard
from solver.SComponent import SComponent

from visualizer.solution_visualizer import visualize_sd, anim_prepare, anim_update_sd, anim_stop
import visualizer.common as v

def _is_facing_up(row: int, col: int) -> bool:
    ret = col % 2 == 0
    if 0 <= row <= 2:
        return ret
    return not ret

class SolDesc:
    def __init__(
        self,
        bd: list[list[int]],
        piece_used: list[bool],
        n_pieces_used: int,
        last_row: int,
        last_col: int
    ) -> None:
        self.bd = bd
        self.piece_used = piece_used
        self.n_pieces_used = n_pieces_used
        self.last_row = last_row
        self.last_col = last_col
    
    def is_solved(self) -> bool:
        return self.n_pieces_used == g_n_pieces
    
    def place(self, r: int, c: int, p: int, t: int) -> list[list[int]] | None:
        """
        Test if we can place the t_th transformation of the p_th piece on r, c
        """
        
        if self.bd[r][c] != -1:
            return None
        
        rd_list, cd_list = [], []
        
        destx = g_comps[p][t].destx
        desty = g_comps[p][t].desty
        
        for i in range(g_comps[p][t].ndest):
            x = g_board.x_map[r][c] + destx[i]
            y = g_board.y_map[r][c] + desty[i]
            rd = g_board.get_row(x, y)
            cd = g_board.get_col(x, y)
            if (rd is None) or (cd is None):
                return None
            if self.bd[rd][cd] != -1:
                return None
            rd_list.append(rd)
            cd_list.append(cd)
        
        new_bd = []
        for row in self.bd:
            new_bd.append(row[:])
        
        new_bd[r][c] = p
        for i in range(g_comps[p][t].ndest):
            rd = rd_list[i]
            cd = cd_list[i]
            new_bd[rd][cd] = p
        
        return new_bd
    
    def iter_next(self) -> Generator["SolDesc", None, None]:
        # Find the first available cell
        r, c = self.last_row, self.last_col
        while not (r == 5 and c == 6):
            if self.bd[r][c] == -1:
                break
            c += 1
            if c == g_row_length[r]:
                r += 1
                c = 0
        
        for piece_index in range(g_n_pieces):
            if self.piece_used[piece_index]:
                continue
            
            for trans_index in range(g_n_transforms[piece_index]):
                if g_comps[piece_index][trans_index].is_facing_up != _is_facing_up(r, c):
                    continue
                
                new_bd = self.place(r, c, piece_index, trans_index)
                if new_bd:
                    pu = self.piece_used[:]
                    pu[piece_index] = True
                    yield SolDesc(new_bd, pu, self.n_pieces_used + 1, r, c)

@dataclass
class Solution:
    sd: SolDesc
    it: int

# ================================================

g_board: SBoard = None # type: ignore
g_comps: list[list[SComponent]] = None # type: ignore
g_n_pieces: int = None # type: ignore
g_n_transforms: list[int] = None # type: ignore
g_row_length: list[int] = None # type: ignore
def _setup(board: SBoard, comps: list[list[SComponent]]):
    global g_board, g_comps
    g_board = board
    g_comps = comps
    
    global g_n_pieces, g_n_transforms
    g_n_pieces = len(comps)
    g_n_transforms = [len(t) for t in comps]
    
    global g_row_length
    g_row_length = [len(row) for row in board.bd]
    
# ================================================

def _solve_all(board: SBoard, comps: list[list[SComponent]]) -> Generator[Solution, None, None]:
    _setup(board, comps)
    
    # Find the first available cell
    r, c = 0, 0
    while not (r == 5 and c == 6):
        if board.bd[r][c] == -1:
            break
        c += 1
        if c == g_row_length[r]:
            r += 1
            c = 0
    
    stack: list[SolDesc] = [
        SolDesc(board.bd, [False for _ in range(g_n_pieces)], 0, r, c)
    ]
    
    it = 0
    while stack:
        sd = stack.pop()
        it += 1
        
        if sd.is_solved():
            yield Solution(sd, it)
            continue
        
        for next_sd in sd.iter_next():
            stack.append(next_sd)
    
def solve(board: SBoard, comps: list[list[SComponent]], solve_all: bool = False) -> list[Solution]:
    ret = []
    for sol in _solve_all(board, comps):
        if not solve_all:
            return [sol]
        ret.append(sol)
    return ret